import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

# 线性回归
def cal_w_b(x, y):
    x_mean = x.sum()/x.size
    fen_mu = 0.0
    xi2 = 0.0
    b = 0.0
    for xi, yi in zip(x, y):
        fen_mu += yi*(xi-x_mean)
        xi2 += xi**2
    fen_zi = xi2 - (x.sum())**2/x.size
    w = fen_mu/fen_zi
    for xi, yi in zip(x, y): 
        b += (yi-w*xi)
    return w, b/x.size

# 预测函数
def y(w, b, x):
    df_predict = pd.DataFrame(x, columns=['面积(平方)'])
    l = []
    for i in x:
        l.append(w*i+b)
    df_predict['预测价格(万)'] = l
    return df_predict

def main():
    #  数据集
    df_data = [
        [77.36, 470],
        [116.74, 730],
        [116.7, 760],
        [100.68, 680],
        [116.1, 700],
        [115.81, 720],
        [104.24, 700],
        [106.73, 690],
        [115.86, 730]
    ]
    df_house = pd.DataFrame(df_data, columns=['面积(平方)', '价格(万)'])
    # 测试数据
    test_date = [56.6, 78.4, 58, 123.5, 56.8, 77, 150.6]
    
    # 预测
    w, b = cal_w_b(df_house['面积(平方)'], df_house['价格(万)'])
    print('w:', w, 'b:', b)
    df_predict = y(w, b, test_date)
    print(df_predict)

    # 可视化
    plt.rcParams['font.sans-serif'] = 'SimHei' 
    plt.rcParams['axes.unicode_minus'] = False
    plt.xlabel('面积(平方)')
    plt.ylabel('价格(万)')
    plt.scatter(df_house['面积(平方)'], df_house['价格(万)'])
    plt.plot(df_predict['面积(平方)'], df_predict['预测价格(万)'])
    plt.legend(['数据集','预测'])
    plt.savefig('线性回归.png')

if __name__ == '__main__':
    main()